import shutil
import re
import os
import yaml
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import plotly.graph_objects as go
import requests
import subprocess
from mpl_toolkits.mplot3d import Axes3D


def setup_modules(module):
    '''Download the modules from gitee.

    Parameters
    ----------
    `module` : str  
        Module name. 'all' for all modules.
        Currently, support 'opensees2yaml' and 'smartanalyze'.
    
    Outputs
    -------
    The specified modules will be downloaded from gitee, saved in `modules` folder
    '''
    os.makedirs('modules', exist_ok=True)
    if module == 'opensees2yaml':
        res = requests.get("https://gitee.com/donghl/codes/pxmrju08vze3b1dohsqal45/raw?blob_name=opensees2yaml.tcl")
        with open('modules/opensees2yaml.tcl', 'w') as f:
            f.write(res.text)
    if module == 'smartanalyze':
        res = requests.get("https://gitee.com/donghl/codes/r7jw1ntbd6uq8gfkoey0c80/raw?blob_name=SmartAnalyze.tcl")
        with open('modules/SmartAnalyze.tcl', 'w') as f:
            f.write(res.text)
    if module == 'all':
        setup_modules('opensees2yaml')
        setup_modules('smartanalyze')


class VisualizerMixin():
    '''The mixin that provides OpsModel class the capability to visualize the model.
    '''
    def view_deformed(self, classify=None, factor=1, show=False, project=0):
        '''Visualize the deformed structure.
        
        Parameters
        ----------
        `classify` : collable, None  
            If is None, the elements are not classified. 
            Else, specify a function that reads the element tag and return the group name.
        
        `factor` : number  
            The factor that is multiplied to the deformation.
        
        `show` : bool  
            Whether show the figure immediately.
        
        `project` : int  
            The projection axis. Only useful in 3D case.
            If ==0, do not project.
            If > 0, the axis to project the 3D structure to 2D space.
        '''
        if self.get_ndm() == 3 and project == 0:
            print('3D plotting interactively is not supported for now.')
            return
        else:
            frames = self.construct_deformed2D(classify=classify, factor=factor, project=project)
            fig = go.Figure(data=frames[0]['data'], layout=self.construct_animation_layout(len(frames)), frames=frames[1:])
        if show:
            fig.show()
        fig.write_html(f'{self.rc["folder_output"]}/{self.name}/{self.rc["filename_plot_deformed"]}')
        return fig

    def view_eigen(self, classify=None, factor=1, show=False, project=0):
        '''Visualize the eigen values.

        Parameters
        ----------
        See `Visualizer_view_deformed`
        '''
        self.load_yaml()
        if not (self.rc and self.rc['plotter_eigen']):
            print('No plotter is specified.')
            return
        if self.rc['plotter_eigen'] == 'plotly':
            if self.get_ndm() == 2 or project != 0:
                frames = self.construct_eigen2D(classify=classify, factor=factor, project=project)
                fig = go.Figure(data=frames[0]['data'], layout=self.construct_animation_layout(len(frames)), frames=frames)
                if show:
                    fig.show()
                fig.write_html(f'{self.rc["folder_output"]}/{self.name}/{self.rc["filename_plot_eigen"]}')
                return fig
            else:
                print('3D plotter for view_eigen is not supported now.')
                return
        else:
            print('Currently only support "plotly" plotter')
            return

    def view_undeformed(self, classify=None, show=False, project=0):
        '''Visualize the undeformed structure.

        Parameters
        ----------
        See `Visualizer.view_deformed`
        '''
        self.load_yaml()
        if not (self.rc and self.rc['plotter_undeformed']):
            print("No plotter is specified.")
            return
        if self.rc['plotter_undeformed'] == 'pyplot':
            if self.get_ndm() == 2 or project != 0:
                fig, ax = plt.subplots(1, 1)
                data = self.construct_lines2D(classify=classify, project=project)
                for sequence in data:
                    ax.plot(sequence['x'], sequence['y'], label=sequence['name'], marker='o', markersize=2)
                if classify is not None:    
                    ax.legend(loc='best', fontsize='small')
            else:
                fig = plt.figure()
                ax = fig.add_subplot(111, projection='3d')
                data = self.construct_lines3D(classify)
                for sequence in data:
                    i = 0
                    ihist = 0
                    while i < len(sequence['x']):
                        while sequence['x'][i] is not None:
                            i += 1
                        x = sequence['x'][ihist:i]
                        y = sequence['y'][ihist:i]
                        z = sequence['z'][ihist:i] 
                        print({'x':x, 'y':y, 'z':z})
                        ax.plot3D(x, y, z, marker='o', markersize=2)
                        i += 1
                        ihist = i
            if show:
                plt.show()
            fig.savefig(f'{self.rc["folder_output"]}/{self.name}/{self.rc["filename_plot_undeformed"]}')
            return fig
        else:
            print('The plotter is not recognized.')
            return

    def load_yaml(self):
        yaml_path = f'{self.rc["folder_output"]}/{self.name}/{self.rc["filename_yaml"]}'
        if not os.path.exists(yaml_path):
            print(f'Yaml file not found: {yaml_path}')
            return
        else:
            with open(yaml_path, 'r') as f:
                self.sdict = yaml.load(f, Loader=yaml.SafeLoader)
            return self.sdict
    
    def get_ndm(self):
        if not hasattr(self, 'sdict'):
            self.load_yaml()
        return self.sdict['ndm']

    def get_ndf(self):
        if not hasattr(self, 'sdict'):
            self.load_yaml()
        return self.sdict['ndf']

    def get_eigen(self):
        if not hasattr(self, 'sdict'):
            self.load_yaml()
        nodetags = list(self.sdict['nodes'].keys())
        eigen = self.sdict['nodes'][nodetags[0]]['eigen']
        if eigen is None:
            return 0
        else:
            return len(eigen)
    

    def construct_nodetags(self, classify=None):
        '''Get the node tags for each element. Then construct the node tags to a list.  

        `classify` : function(int) -> str  
            The classification function. Read the element tag, and return a the group name.
        
        Returns
        -------
        `nodetags` : List<int>  
            A list of node tags.
        '''
        eletags = self.sdict['elements'].keys()
        if classify is None:
            grouping = {'': eletags}
        else:
            grouping = {}
            for eletag in eletags:
                group = classify(eletag)
                if group not in grouping:
                    grouping[group] = []
                grouping[group].append(eletag)
        def ele2node(eletag):
            eles = self.sdict['elements']
            res = []
            if len(eles[eletag]['connections']) == 2:
                connections = eles[eletag]['connections']
            else:
                connections = eles[eletag]['connections'] + [eles[eletag]['connections'][0]]
            for connection in connections:
                res.append(connection)
            res.append(None)
            return res
        def eles2nodes(eletags):
            res = []
            for eletag in eletags:
                res += ele2node(eletag)
            return res
        res = {}
        for group, eletags in grouping.items():
            res[group] = eles2nodes(eletags)
        return res

    def construct_lines2D(self, classify=None, project=0):
        if project == 1:
            lx, ly = 1, 2
        elif project == 2:
            lx, ly = 0, 2
        else:
            lx, ly = 0, 1

        nodetags_grouping = self.construct_nodetags(classify)
        res = []
        for group, nodetags in nodetags_grouping.items():    
            groupdict = {}
            groupdict['x'] = [self.sdict['nodes'][t]['coords'][lx] if t is not None else None for t in nodetags]
            groupdict['y'] = [self.sdict['nodes'][t]['coords'][ly] if t is not None else None for t in nodetags]
            groupdict['name'] = group
            res.append(groupdict)
        return res
    
    def construct_lines3D(self, classify=None):
        nodetags_grouping = self.construct_nodetags(classify)
        res = []
        for group, nodetags in nodetags_grouping.items():    
            groupdict = {}
            groupdict['x'] = [self.sdict['nodes'][t]['coords'][0] if t is not None else None for t in nodetags]
            groupdict['y'] = [self.sdict['nodes'][t]['coords'][1] if t is not None else None for t in nodetags]
            groupdict['z'] = [self.sdict['nodes'][t]['coords'][2] if t is not None else None for t in nodetags]
            groupdict['name'] = group
            res.append(groupdict)
        return res
    
    def read_allnodes(self):
        filepath = f'{self.rc["folder_output"]}/{self.name}/{self.rc["filename_allnodes"]}'
        if not os.path.exists(filepath):
            print(f'File do not exist: {filepath}')
            return
        allnodes = np.loadtxt(filepath)
        if not hasattr(self, 'sdict'):
            self.load_yaml()
        if allnodes.shape[1] != len(self.sdict["nodes"].keys()) * 2 and allnodes.shape[1] != len(self.sdict["nodes"].keys()) * 3:
            print(f'The node recorder file shape {allnodes.shape} do not match the amount of nodes {len(self.sdict["nodes"].keys())}.')
            return
        else: 
            self.allnodes = allnodes
            return allnodes

    def construct_deformed2D(self, classify=None, factor=1, project=0):
        if project == 1:
            lx, ly = 1, 2
        elif project == 2:
            lx, ly = 0, 2
        else:
            lx, ly = 0, 1
        self.read_allnodes()
        ix, iy = dict(), dict()
        for i, nodetag in enumerate(self.sdict['nodes'].keys()):
            if project == 0:
                ix[nodetag] = i * 2
                iy[nodetag] = i * 2 + 1
            elif project == 1:
                ix[nodetag] = i * 3 + 1
                iy[nodetag] = i * 3 + 2
            elif project == 2:
                ix[nodetag] = i * 3
                iy[nodetag] = i * 3 + 2
            elif project == 3:
                ix[nodetag] = i * 3
                iy[nodetag] = i * 3 + 1
        nodetags_grouping = self.construct_nodetags(classify)
        frame0 = {"data": self.construct_lines2D(classify), "name": 0}
        frames = [frame0]
        for i in range(self.allnodes.shape[0]):
            data = []
            for name, nodetags in nodetags_grouping.items():
                sequence = {}
                sequence['name'] = name
                sequence['x'] = [self.sdict['nodes'][t]['coords'][lx] + factor * self.allnodes[i, ix[t]] if t is not None else None for t in nodetags]
                sequence['y'] = [self.sdict['nodes'][t]['coords'][ly] + factor * self.allnodes[i, iy[t]] if t is not None else None for t in nodetags]
                sequence['ids'] = [str(i) for i in range(len(nodetags))]
                data.append(sequence)
            frames.append({"data": data, "name":i+1})
        return frames

    def construct_deformed3D(self, classify=None, factor=1):
        self.read_allnodes()
        index = dict()
        for i, nodetag in enumerate(self.sdict['nodes'].keys()):
            index[nodetag] = i * 2
        nodetags_grouping = self.construct_nodetags(classify)
        frame0 = {"data": self.construct_lines3D(classify), "name": 0}
        frames = [frame0]
        for i in range(self.allnodes.shape[0]):
            data = []
            for name, nodetags in nodetags_grouping.items():
                sequence = {}
                sequence['name'] = name
                sequence['x'] = [self.sdict['nodes'][t]['coords'][0] + factor * self.allnodes[i, index[t]  ] if t is not None else None for t in nodetags]
                sequence['y'] = [self.sdict['nodes'][t]['coords'][1] + factor * self.allnodes[i, index[t]+1] if t is not None else None for t in nodetags]
                sequence['z'] = [self.sdict['nodes'][t]['coords'][2] + factor * self.allnodes[i, index[t]+2] if t is not None else None for t in nodetags]
                sequence['ids'] = [str(i) for i in range(len(nodetags))]
                data.append(sequence)
            frames.append({"data": data, "name":i+1})
        return frames
    
    def construct_eigen2D(self, classify=None, factor=1, project=0):
        if project == 1:
            lx, ly = 1, 2
        elif project == 2:
            lx, ly = 0, 2
        else:
            lx, ly = 0, 1
        eigen_count = self.get_eigen()
        if not eigen_count:
            print('No eigen is recorded in opensees2yaml.')
            return
        # estimate the factor.
        EIG_MODEL_RATIO = 0.05
        xmin, xmax, ymin, ymax, exmax, eymax = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
        for _, node in self.sdict["nodes"].items():
            if xmin > node["coords"][lx]:
                xmin = node["coords"][lx]
            if xmax < node["coords"][lx]:
                xmax = node["coords"][lx]
            if ymin > node["coords"][ly]:
                ymin = node["coords"][ly]
            if ymax < node["coords"][ly]:
                ymax = node["coords"][ly]
            if exmax < abs(node["eigen"][0][lx]):
                exmax = abs(node["eigen"][0][lx])
            if eymax < abs(node["eigen"][0][ly]):
                eymax = abs(node["eigen"][0][ly])
        fact1 = (xmax - xmin) * EIG_MODEL_RATIO / exmax
        fact2 = (ymax - ymin) * EIG_MODEL_RATIO / eymax
        factor = factor * min(fact1, fact2)
        # construct frames
        nodetags_grouping = self.construct_nodetags(classify)
        frame0 = {"data": self.construct_lines2D(classify), "name": 0}
        frames = [frame0]
        for i in range(eigen_count):
            data = []
            for name, nodetags in nodetags_grouping.items():
                sequence = {}
                sequence['name'] = name
                sequence['x'] = [self.sdict['nodes'][t]['coords'][lx] + factor * self.sdict['nodes'][t]['eigen'][i][lx] if t is not None else None for t in nodetags]
                sequence['y'] = [self.sdict['nodes'][t]['coords'][ly] + factor * self.sdict['nodes'][t]['eigen'][i][ly] if t is not None else None for t in nodetags]
                sequence['ids'] = [str(i) for i in range(len(nodetags))]
                data.append(sequence)
            frames.append({"data": data, "name":i+1})
        return frames

    
    def construct_animation_layout(self, frames_count):
        layout = {}
        layout['updatemenus'] = [
            {
                'buttons': [{
                    'args': [None, {'frame': {'duration': 500, 'redraw': False},
                                    'fromcurrent': True,
                                    'transition': {'duration': 300, 'easing': 'quadric-in-out'}}],
                    'label': 'Play',
                    'method': 'animate'
                }, {
                    'args': [[None], {'frame': {'duration': 0, 'redraw': False},
                                    'mode': 'immediate',
                                    'transition': {'duration': 0}}],
                    'label': 'Pause',
                    'method': 'animate'
                }],
                'direction': 'left',
                'pad': {'r': 10, 't': 80},
                'type': 'buttons',
                'showactive': False,
                'x': 0.1,
                'xanchor': 'right',
                'y': 0,
                'yanchor': 'top'
            }
        ]
        layout['sliders'] = [{
            'active': 0,
            'yanchor': 'top',
            'xanchor': 'left',
            'transition': {'duration': 300, 'easing': 'cubic-in-out'},
            'pad': {'b': 10, 't': 50},
            'len': 0.9,
            'x': 0.1,
            'y': 0,
            'steps': [{
                'args': [[i], {'frame': {'duration': 300, 'redraw': False},
                               'mode': 'immediate',
                               'transition': {'duration': 300}}],
                'label': i,
                'method': 'animate'
            } for i in range(frames_count)]
        }]
        return layout

class OpsModel(VisualizerMixin):
    '''This class represents an OpenSees model.
    With the help of this class, users can easily deal with the model scripts,
    outputs, data, and plots.
    The script workspace root is the current folder.
    The output folder is self.folder.
    '''

    _rc_default = dict(
        filename_model='model.tcl',
        filename_plot_undeformed='model.png',
        filename_plot_deformed='interactive.html',
        filename_plot_eigen='eigen.html',
        filename_yaml='model.yml',
        filename_allnodes='all.nodes',
        filename_opslog='ops.log',
        filename_data='data.csv',
        filename_script='script.tcl',
        folder_output='.out',
        folder_archive_data='data',
        folder_archive_image='image',
        extension_output='out',
        archive_data_regexes=['data.csv'],
        archive_image_regexes=['model.png', 'eigen.html', 'interactive.html'],
        save_cmd=True,
        save_undeformed=True,
        save_deformed=True,
        plotter_undeformed='pyplot',
        plotter_eigen='plotly',
        plotter_deformed='plotly',
    )
    '''rc filenames for the plots.'''

    def __init__(self, name, **kwargs):
        '''
        Parameters
        ----------
        `name` : str  
            Every OpsModel instance must has a name. The name is used in the following contexts:  
                1. All the outputs are stored in the {folder_output}/{name} folder.
                2. All the data and images are archived to the archive folder with filename {name.xxx}

        `kwargs` : dict  
            A dict that changes the default values of `OpsModel._rc_default`
        '''
        self.name = name
        '''See `OpsModel`'''
        rc = self._rc_default.copy()
        rc.update(kwargs)
        self.rc = rc
        '''See `OpsModel`'''
        self.folder = f'{self.rc["folder_output"]}/{self.name}'
        '''The operating folder'''
        os.makedirs(self.folder, exist_ok=True)
        self.script = ''
        '''str. The tcl script of the structure.'''
        self.cmd = ''
        '''str. The shell command to run the tcl file.
        If cmd == '', then run the script interactively.'''
        self.data = None
        '''pandas.DataFrame. The collected data.'''
    
    def run_cmd(self):
        '''Run the OpenSees model by cmd.
        
        Returns
        -------
        `stderr` : str  
            The output of opensees.
        '''
        res = subprocess.run(self.cmd.split(' '), text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        with open(f'{self.folder}/{self.rc["filename_opslog"]}', 'w') as f:
            f.write(res.stderr)
            f.write(res.stdout)
        return res.stderr
    
    def run_script(self, write=True):
        '''Run the OpenSees model by script.

        Parameters
        ----------
        `write` : bool  
            If True, the script is written to the output folder.
            Else the script is run in the console and not written.
        
        Returns
        -------
        `stderr` : str  
            The output of opensees.
        '''
        if write:
            with open(f'{self.folder}/{self.rc["filename_script"]}', 'w') as f:
                f.write(self.script)
            if not self.cmd:
                self.cmd = f'opensees {self.folder}/{self.rc["filename_script"]}'
                stderr = self.run_cmd()
        else:
            res = subprocess.run('opensees', input=self.script, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            with open(f'{self.folder}/{self.rc["filename_opslog"]}', 'w') as f:
                f.write(res.stderr)
                f.write(res.stdout)
            stderr = res.stderr
        return stderr

    
    def collect(self):
        '''Collet the output data from the recorders. This is a virtual class.
        This method must be overwritten in the child classes.
        
        Parameters
        ----------
        `write` : bool  
            Whether write the csv file to the output folder.
            Only set write=False when this method is overwritten.
        
        Outputs
        -------
        self.data is changed.
        
        Returns
        -------
        `data` : `pandas.DataFrame`.
        '''
        res = dict()
        filenames = os.listdir(self.folder)
        for filename in filenames:
            if filename.endswith(self.rc["extension_output"]):
                data = np.loadtxt(f'{self.folder}/{filename}')
                for i in range(data.shape[1]):
                    colname = f'{filename[:-len(self.rc["extension_output"])-1]}_col{i}'
                    res[colname] = data[:, i]
        self.data = pd.DataFrame(res)
        return self.data

    def write_data(self):
        '''Write the collected data into a csv file.
        '''
        self.data.to_csv(f'{self.folder}/{self.rc["filename_data"]}')

    def read_data(self):
        '''Read the csv file generated by `OpsModel.write_data`.
        '''
        self.data = pd.read_csv(f'{self.folder}/{self.rc["filename_data"]}', index_col=0)
        return self.data

    def archive_data(self):
        '''Collect the data to the {rc.folder_archive_data} folder.
        The data regex is defined in a list in rc.data_regexes
        
        Output
        ------
        Files matching any of the rc["archive_data_regexes"] will be copied to the folder rc["folder_archive_image"].  
        '''
        if not self.rc['archive_data_regexes']:
            print('No file is to be archived')
            return
        pattern = re.compile('|'.join([f'({s})' for s in self.rc['archive_data_regexes']]))
        existing_files = os.listdir(f'{self.rc["folder_output"]}/{self.name}')
        copy_files = [s for s in existing_files if pattern.match(s) is not None]
        if not copy_files:
            print('No file matches the regex')
            return
        os.makedirs(self.rc['folder_archive_data'], exist_ok=True)
        for filename in copy_files:
            shutil.copy(f'{self.rc["folder_output"]}/{self.name}/{filename}', f'{self.rc["folder_archive_data"]}/{self.name}_{filename}')
            print(f'Copied {self.rc["folder_output"]}/{self.name}/{filename}')
    
    def archive_image(self):
        '''Collect the images to the {rc.folder_archive_image} folder.
        The images regex is defined in a list in rc.image_regexes
        
        Output
        ------
        Files matching any of the rc["archive_image_regexes"] will be copied to the folder rc["folder_archive_image"].  
        '''
        if not self.rc['archive_image_regexes']: 
            print('No file is to be archived.')
            return
        pattern = re.compile('|'.join([f'({s})' for s in self.rc['archive_image_regexes']]))
        existing_files = os.listdir(f'{self.rc["folder_output"]}/{self.name}')
        copy_files = [s for s in existing_files if pattern.match(s) is not None]
        if not copy_files:
            print('No file matches rc folder_image_regexes')
            return
        os.makedirs(self.rc['folder_archive_image'], exist_ok=True)
        for filename in copy_files:
            shutil.copy(f'{self.rc["folder_output"]}/{self.name}/{filename}', f'{self.rc["folder_archive_image"]}/{self.name}_{filename}')
            print(f'Copied {self.rc["folder_output"]}/{self.name}/{filename}')
    

sample_script_2D = '''source modules/SmartAnalyze.tcl
source modules/opensees2yaml.tcl
wipe
wipeAnalysis
model BasicBuilder -ndm 2 -ndf 3
node 0 0. 0.
node 1 2000. 0.
node 2 0. 1000.
node 3 2000. 1000.
fix 0 1 1 1
fix 1 1 1 1
mass 2 10. 10. 10.
mass 3 10. 10. 10.
geomTransf Linear 1
element elasticBeamColumn 1 0 2 1000. 10000. 1000000. 1
element elasticBeamColumn 2 1 3 1000. 10000. 1000000. 1
element elasticBeamColumn 3 2 3 1000. 10000. 1000000. 1
pattern Plain 1 Linear {
    load 2 1000. 0. 0.
}
file mkdir .out
file mkdir .out/sample_2D
recorder Node -file .out/sample_2D/disp_story0.out -time -node 2 -dof 1 disp
recorder Element -file .out/sample_2D/force_story00.out -ele 1 force
recorder Element -file .out/sample_2D/force_story01.out -ele 2 force
recorder Node -file .out/sample_2D/all.nodes -nodeRange 0 10000 -dof 1 2 disp
set yaml [opensees2yaml 1]
set f [open ".out/sample_2D/model.yml" "w"]
puts $f $yaml
close $f
constraints Plain
numberer Plain
system BandGeneral
test EnergyIncr 1.0e-6 200
algorithm Newton
integrator DisplacementControl 2 1 1
analysis Static
analyze 10
'''

sample_script_3D = '''source modules/SmartAnalyze.tcl
source modules/opensees2yaml.tcl
wipe
wipeAnalysis
model BasicBuilder -ndm 3 -ndf 6
node 0 0. 0. 0.
node 1 2000. 1000. 0.
node 2 0. 0. 1000
node 3 2000. 1000. 1000
fix 0 1 1 1 1 1 1
fix 1 1 1 1 1 1 1
mass 2 10. 10. 10. 0. 0. 0.
mass 3 10. 10. 10. 0. 0. 0.
geomTransf Linear 1 0. 1. 1.
element elasticBeamColumn 1 0 2 1000. 10000. 200 1.0e9 1.0e5 1.0e5 1
element elasticBeamColumn 2 1 3 1000. 10000. 200 1.0e9 1.0e5 1.0e5 1
element elasticBeamColumn 3 2 3 1000. 10000. 200 1.0e9 1.0e5 1.0e5 1
pattern Plain 1 Linear {
    load 2 1000. 0. 0. 0. 0. 0.
}
file mkdir .out
file mkdir .out/sample_3D
recorder Node -file .out/sample_3D/disp_story0.out -time -node 2 -dof 1 2 disp
recorder Element -file .out/sample_3D/force_story00.out -ele 1 force
recorder Element -file .out/sample_3D/force_story01.out -ele 2 force
recorder Node -file .out/sample_3D/all.nodes -nodeRange 0 10000 -dof 1 2 3 disp
set yaml [opensees2yaml 1]
set f [open ".out/sample_3D/model.yml" "w"]
puts $f $yaml
close $f
constraints Plain
numberer Plain
system BandGeneral
test EnergyIncr 1.0e-6 200
algorithm Newton
integrator DisplacementControl 2 1 1
analysis Static
analyze 10
'''

class MultiStoryMixin():
    def plot_profile(self, label):
        df = self.data.groupby('Story').max().reset_index()
        x = df[label].to_numpy()
        x = np.hstack([0, x])
        y = df['Story'].to_numpy()
        y = np.hstack([0, y])
        plt.figure()
        plt.plot(x, y)
        plt.xlabel(label)
        plt.ylabel('Story')
        plt.show()

__all__ = [
    'setup_modules',
    'VisualizerMixin',
    'MultiStoryMixin',
    'OpsModel',
    'sample_script_2D',
    'sample_script_3D',
]